package com.shenjies88.practice.baseframework.controller;

import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletResponse;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

/**
 * @author by shenjies88
 * @since 2025/9/13 17:32
 */
@Slf4j
@AllArgsConstructor
@RequestMapping("/chatModel")
@RestController
public class ChatModelController {

    @Resource(name = "qwen3ChatModel")
    private final OllamaChatModel qwen3ChatModel;

    @Resource(name = "dsr1ChatModel")
    private final OllamaChatModel dsr1ChatModel;

    @GetMapping("/chat")
    public String chat(@RequestParam String msg) {
        ChatResponse chatResponse = dsr1ChatModel.call(new Prompt(msg));
        return chatResponse.getResult().getOutput().getText();
    }

    @GetMapping("/streamChat")
    public Flux<String> streamChat(@RequestParam String msg, HttpServletResponse response) {
        response.setCharacterEncoding("UTF-8");
        Flux<ChatResponse> stream = dsr1ChatModel.stream(new Prompt(msg));
        return stream.mapNotNull(resp -> {
            String text = resp.getResult().getOutput().getText();
            log.info(text);
            return text;
        });
    }

}
